import os
import pcse
import matplotlib.pyplot as plt
from pcse.input import CABOFileReader,YAMLAgroManagementReader,ExcelWeatherDataProvider,NASAPowerWeatherDataProvider
#from pcse.util import WOFOST72SiteDataProvider
from pcse.base import ParameterProvider
from pcse.models import Wofost72_WLP_CWB

#define the data directory
data_dir = r'D:\\pythonproject\\pcse_walkthrough\\userdata\\pcse_examples'


#Crop parameters
"""
The crop parameters consist of parameter names 
and the corresponding parameter values that are needed to parameterize the components of the crop simulation model. 
These are crop-specific values regarding phenology, assimilation, respiration, biomass partitioning, etc.
The parameter file for sugar beet is taken from the crop files in the WOFOST Control Centre.

作物参数包括参数名称和用于参数化作物模拟模型组件所需的相应参数值。
这些是与作物相关的物候、吸收、呼吸、生物量分配等参数值。糖用甜菜的参数文件是从WOFOST控制中心的作物文件中提取的。
"""
cropfile = os.path.join(data_dir+ "\\quickstart_part2", 'sug0601.crop')
cropdata = CABOFileReader(cropfile)
#print(cropdata)

#Soil parameters
"""
The soildata dictionary provides the parameter name/value pairs related to the soil type and soil physical properties.
The number of parameters is variable depending on the soil water balance type that is used for the simulation. For this example, 
we will use the water balance for freely draining soils and use the soil file for medium fine sand: ec3.soil. 
This file is also taken from the soil files in the WOFOST Control Centre
土壤数据字典提供了与土壤类型和土壤物理性质相关的参数名称/值对。参数的数量根据用于模拟的土壤水分平衡类型而变化。
对于这个例子，我们将使用水平衡来自由排水土壤，并使用土壤锉来处理中细砂：ec3.soil。该文件也取自WOFOST控制中心的土壤文件
"""
soilfile = os.path.join(data_dir+ "\\quickstart_part2", 'ec3.soil')
soildata = CABOFileReader(soilfile)

#Site parameters
"""
现场参数提供了与作物或土壤无关的辅助参数。例如，水平衡的初始条件，如初始土壤含水量（WAV）和初始和最大地表蓄水量（SSI、SSMAX）。
大气中的二氧化碳浓度也是一个典型的现场参数。目前，我们可以直接在Python命令行上将这些参数定义为一个简单的Python字典。
但是，使用WOFOST72SiteDataProvider更方便，它记录了站点参数并提供了合理的默认值
"""
sitedata =  {'SMLIM': 0.4, 'NOTINF': 0, 'CO2': 360.0, 'SSI': 0.0, 'SSMAX': 0.0, 'IFUNRN': 0, 'WAV': 100.0}
#print(sitedata)

parameters = ParameterProvider(cropdata=cropdata, soildata=soildata, sitedata=sitedata)

"""
农业管理输入提供了农业活动的开始日期、作物模拟的start_date/start_type、作物仿真的end_date/end_type和作物模拟的最大持续时间。
包括后者是为了避免不切实际的长时间模拟，例如由于过高的温度总和要求。

农业管理输入是用一种称为YAML的特殊语法定义的，它允许轻松创建定义农业管理所需的更复杂的结构。
Wageningen sugarbeet_calendar.agro 中甜菜的农业管理文件可以使用YAMLAgroManagementReader读取:
"""
agromanagement_file = os.path.join(data_dir+ "\\quickstart_part2", 'sugarbeet_calendar.agro')
agromanagement = YAMLAgroManagementReader(agromanagement_file)
print(agromanagement)

"""
运行模拟需要每日天气变量。PCSE中有几个数据提供程序用于读取天气数据, 请参阅天气数据提供程序部分以获取概述。
"""
wdp = ExcelWeatherDataProvider(xls_fname=data_dir+ '\\quickstart_part3\\nl1.xlsx',force_reload=True)
#wdp = NASAPowerWeatherDataProvider(latitude=51.97, longitude=5.67)
print(wdp)

#Importing, initializing and running a PCSE model
wofsim = Wofost72_WLP_CWB(parameters, wdp, agromanagement)

"""
We can then run the simulation and show some final results such as the anthesis and harvest dates (DOA, DOH), total biomass (TAGP) and maximum LAI (LAIMAX).
Next, we retrieve the time series of daily simulation output using the get_output() method on the WOFOST object:
"""
wofsim.run_till_terminate()
output = wofsim.get_output()
#print(output)
"""
As the output is returned as a list of dictionaries, we need to unpack these variables from the list of output:

Finally, we can generate some figures of WOFOST variables such as the development (DVS), total biomass (TAGP), 
leaf area index (LAI) and root-zone soil moisture (SM) using the MatPlotLib plotting package:

"""
varnames = ["day", "DVS", "TAGP", "LAI", "SM", "TWSO"]
tmp = {}
for var in varnames:
    tmp[var] = [t[var] for t in output]

day = tmp.pop("day")
fig, axes = plt.subplots(nrows=2, ncols=3, figsize=(15,8))
for var, ax in zip(["DVS", "TAGP", "LAI", "SM", "TWSO"], axes.flatten()):
    ax.plot(day, tmp[var], 'b-')
    ax.set_title(var)
fig.autofmt_xdate()
fig.savefig('sugarbeet.png')